from sympy import false

from playLA._globa import is_zero
from playLA.Matrix import Matrix
from playLA.Vector import Vector


class LinearSystem:
    def __init__(self, A, b):
        assert A.row_num() == len(b), "row number of A must be equal to the length of b"
        self._m = A.row_num()
        self._n = A.col_num()
        # assert self._m == self._n  # TODO: no this restriction

        self.Ab = [Vector(A.row_vector(i).underlying_list() + [b[i]]) for i in range(self._m)]
        # 每一行主元所在的列
        self.pivots = []

    def _max_row(self, index_i, index_j, n):
        best, ret = self.Ab[index_i][index_j], index_i
        for i in range(index_i + 1, n):
            if self.Ab[i][index_j] < best:
                best, ret = self.Ab[i][index_j], i
        return ret

    def _forward(self):
        # 前向高斯消元
        i, k = 0, 0
        while i < self._m and k < self._n:
            # 看Ab[i][k]位置是否可以是主元
            # 不管主元是不是0都把主元位置最大的放前面
            max_row = self._max_row(i, k, self._m)
            self.Ab[i], self.Ab[max_row] = self.Ab[max_row], self.Ab[i]

            if is_zero(self.Ab[i][k]):
                k += 1
            else:
                # 主元归一
                self.Ab[i] = self.Ab[i] / self.Ab[i][i]
                # 把主元列 下面的位置变为 0
                for j in range(i + 1, self._m):
                    # 第 i 行主元已经为1了，乘以j行i位置的系数再做减法消元
                    self.Ab[j] = self.Ab[j] - self.Ab[j][i] * self.Ab[i]
                self.pivots.append(k)
                i += 1

    def _backward(self):
        # 后向把除了主元以外的位置都归0
        # 从后往前
        n = len(self.pivots)
        for i in range(n - 1, -1, -1):
            k = self.pivots[i]
            # Ab[i][k] 为主元
            for j in range(i - 1, -1, -1):
                self.Ab[j] = self.Ab[j] - self.Ab[j][k] * self.Ab[i]

    def gauss_jordan_elimination(self):
        # 前向把所有主元归一，主对角线上都为1
        self._forward()
        # 后向把除了主元以外的都归0
        self._backward()

        for i in range(len(self.pivots), self._m):
            if not is_zero(self.Ab[i][-1]):
                return False
        return True

    def fancy_print(self):
        for i in range(self._m):
            print(" ".join(str(self.Ab[i][j]) for j in range(self._n)), end=" ")
            print("|", self.Ab[i][-1])
